package join;

import lombok.extern.slf4j.Slf4j;
import org.apache.lucene.analysis.core.WhitespaceAnalyzer;
import org.apache.lucene.document.*;
import org.apache.lucene.index.*;
import org.apache.lucene.search.*;
import org.apache.lucene.search.join.*;
import org.apache.lucene.search.join.ScoreMode;
import org.apache.lucene.store.MMapDirectory;
import org.apache.lucene.tests.search.CheckHits;
import org.apache.lucene.util.BitSet;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import java.io.IOException;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;

import static org.junit.jupiter.api.Assertions.assertEquals;

/**
 test the parent-child to excute the join-parent-child-query
 */
@Slf4j
public class JoinDemo {
    MMapDirectory mMapDirectory = null;
    IndexWriter w =null;
    Random random = null;

    @BeforeEach
    public void setUp(){

        String uri = "F:\\integration\\search\\lucene-data\\join";
        try {
            mMapDirectory = new MMapDirectory(Paths.get(uri));
            IndexWriterConfig indexWriterConfig = new IndexWriterConfig();
            indexWriterConfig.setOpenMode(IndexWriterConfig.OpenMode.CREATE);

            WhitespaceAnalyzer whitespaceAnalyzer = new WhitespaceAnalyzer(30);
            w = new IndexWriter(mMapDirectory, new IndexWriterConfig(whitespaceAnalyzer));
             random = new Random();

        } catch (IOException e) {
            throw new RuntimeException("mmap failed");
        }

    }
    @AfterEach
    public void destroyDown(){
        try {
            mMapDirectory.close();
        } catch (IOException e) {
            e.printStackTrace();
        }

    }
    // One resume...
    private Document makeResume(String name, String country) {
        Document resume = new Document();
        //
        resume.add(new StringField("docType", "resume", Field.Store.NO));
        resume.add(new StringField("name", name, Field.Store.YES));
        resume.add(new StringField("country", country, Field.Store.NO));
        return resume;
    }

    // ... has multiple jobs
    private Document makeJob(String skill, int year) {
        Document job = new Document();
        job.add(new StringField("skill", skill, Field.Store.YES));
        job.add(new IntPoint("year", year));
        job.add(new StoredField("year", year));
        return job;
    }

    // ... has multiple qualifications
    private Document makeQualification(String qualification, int year) {
        Document job = new Document();
        job.add(new StringField("qualification", qualification, Field.Store.YES));
        job.add(new IntPoint("year", year));
        job.add(new StoredField("year", year));
        return job;
    }

    private Document getParentDoc(IndexReader reader, BitSetProducer parents, int childDocID)
            throws IOException {
        final List<LeafReaderContext> leaves = reader.leaves();
        final int subIndex = ReaderUtil.subIndex(childDocID, leaves);
        final LeafReaderContext leaf = leaves.get(subIndex);
        final BitSet bits = parents.getBitSet(leaf);
        return leaf.reader().document(bits.nextSetBit(childDocID - leaf.docBase));
    }

    @Test
    public void index() throws Exception {

        final List<Document> docs = new ArrayList<>();

        docs.add(makeJob("java", 2007));
        docs.add(makeJob("python", 2010));
        docs.add(makeResume("Lisa", "United Kingdom"));// dont store the country.
        w.addDocuments(docs);

        docs.clear();
        docs.add(makeJob("ruby", 2005));
        docs.add(makeJob("java", 2006));
        docs.add(makeResume("Frank", "United States"));
        w.addDocuments(docs);

        w.close();

    }
    @Test
    public void search() throws IOException {
        DirectoryReader r = DirectoryReader.open(mMapDirectory);
        IndexSearcher s = new IndexSearcher(r);

        // Create a filter that defines "parent" documents in the index - in this case resumes
        BitSetProducer parentsFilter =
                new QueryBitSetProducer(new TermQuery(new Term("docType", "resume")));
        CheckJoinIndex.check(r, parentsFilter);

        // Define child document criteria (finds an example of relevant work experience)
        BooleanQuery.Builder childQuery = new BooleanQuery.Builder();
        childQuery.add(new BooleanClause(new TermQuery(new Term("skill", "java")), BooleanClause.Occur.MUST));
        childQuery.add(new BooleanClause(IntPoint.newRangeQuery("year", 2006, 2011), BooleanClause.Occur.MUST));

        // Define parent document criteria (find a resident in the UK)
        Query parentQuery = new TermQuery(new Term("country", "United Kingdom"));

        // Wrap the child document query to 'join' any matches
        // up to corresponding parent:
        ToParentBlockJoinQuery childJoinQuery =
                new ToParentBlockJoinQuery(childQuery.build(), parentsFilter, ScoreMode.Avg);

        // Combine the parent and nested child queries into a single query for a candidate
        BooleanQuery.Builder fullQuery = new BooleanQuery.Builder();
        fullQuery.add(new BooleanClause(parentQuery, BooleanClause.Occur.MUST));
        fullQuery.add(new BooleanClause(childJoinQuery, BooleanClause.Occur.MUST));

//        CheckHits.checkHitCollector(random, fullQuery.build(), "country", s, new int[] {2});

        TopDocs topDocs = s.search(fullQuery.build(), 1);

        // assertEquals(1, results.totalHitCount);
        assertEquals(1, topDocs.totalHits.value);
        Document parentDoc = s.doc(topDocs.scoreDocs[0].doc);
        log.info("Parent hit's score is the average of all child scores is:{}",topDocs.scoreDocs[0].score);;
        assertEquals("Lisa", parentDoc.get("name"));

        ParentChildrenBlockJoinQuery childrenQuery =
                new ParentChildrenBlockJoinQuery(
                        parentsFilter, childQuery.build(), topDocs.scoreDocs[0].doc); // id=1
        TopDocs matchingChildren = s.search(childrenQuery, 1);
        assertEquals(1, matchingChildren.totalHits.value);
        assertEquals("java", s.doc(matchingChildren.scoreDocs[0].doc).get("skill"));

        // System.out.println("TEST: now test up");

        // Now join "up" (map parent hits to child docs) instead...:
        ToChildBlockJoinQuery parentJoinQuery = new ToChildBlockJoinQuery(parentQuery, parentsFilter);
        BooleanQuery.Builder fullChildQuery = new BooleanQuery.Builder();
        fullChildQuery.add(new BooleanClause(parentJoinQuery, BooleanClause.Occur.MUST));
        fullChildQuery.add(new BooleanClause(childQuery.build(), BooleanClause.Occur.MUST));

        // System.out.println("FULL: " + fullChildQuery);
        TopDocs hits = s.search(fullChildQuery.build(), 10);
        assertEquals(1, hits.totalHits.value);
        Document childDoc = s.doc(hits.scoreDocs[0].doc);
        // System.out.println("CHILD = " + childDoc + " docID=" + hits.scoreDocs[0].doc);
        assertEquals("java", childDoc.get("skill"));
        assertEquals(2007, childDoc.getField("year").numericValue());
        assertEquals("Lisa", getParentDoc(r, parentsFilter, hits.scoreDocs[0].doc).get("name"));

        // Test with filter on child docs:
        fullChildQuery.add(new TermQuery(new Term("skill", "foosball")), BooleanClause.Occur.FILTER);
        assertEquals(0, s.count(fullChildQuery.build()));

        r.close();
    }
}
